"""
向量化模块
使用Sentence Transformers将文本转为向量
"""
import numpy as np
from typing import List, Dict
from sentence_transformers import SentenceTransformer
import logging
import pickle

from config import EMBEDDING_MODEL_NAME, MODELS_DIR

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class TextEmbedder:
    """文本向量化器"""
    
    def __init__(self, model_name: str = EMBEDDING_MODEL_NAME):
        """
        初始化向量化模型
        
        Args:
            model_name: 模型名称或路径
        """
        logger.info(f"加载向量化模型: {model_name}")
        self.model = SentenceTransformer(model_name)
        self.model_name = model_name
    
    def encode_texts(self, texts: List[str], batch_size: int = 32) -> np.ndarray:
        """
        将文本列表转为向量矩阵
        
        Args:
            texts: 文本列表
            batch_size: 批处理大小
            
        Returns:
            向量矩阵 (n_texts, embedding_dim)
        """
        embeddings = self.model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=True,
            convert_to_numpy=True
        )
        return embeddings
    
    def encode_papers(self, papers: List[Dict]) -> np.ndarray:
        """
        将论文转为向量
        
        Args:
            papers: 论文列表
            
        Returns:
            向量矩阵
        """
        # 组合标题和摘要作为输入文本
        texts = []
        for paper in papers:
            title = paper.get("title", "")
            abstract = paper.get("abstract", "")
            combined_text = f"{title}. {abstract}".strip()
            texts.append(combined_text)
        
        return self.encode_texts(texts)
    
    def save_model(self, save_path: str):
        """保存模型"""
        self.model.save(save_path)
        logger.info(f"模型已保存到: {save_path}")
    
    def load_model(self, load_path: str):
        """加载模型"""
        self.model = SentenceTransformer(load_path)
        logger.info(f"模型已从 {load_path} 加载")


if __name__ == "__main__":
    # 测试向量化
    embedder = TextEmbedder()
    
    test_texts = [
        "COVID-19 is a respiratory disease caused by SARS-CoV-2.",
        "Machine learning improves drug discovery in pharmaceutical research."
    ]
    
    embeddings = embedder.encode_texts(test_texts)
    print(f"向量形状: {embeddings.shape}")
    print(f"前5个维度: {embeddings[0][:5]}")
